import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
import os
import logging
import cv2

# plt.rcParams['font.sans-serif'] = ['SimHei']  # Or other Chinese font
# plt.rcParams['axes.unicode_minus'] = False

class Visualize:
    def __init__(self,log_file="./logs/calibration_param_log.txt") -> None:
        # 设置日志
        os.makedirs(os.path.dirname(log_file), exist_ok=True)
        self.logger = logging.getLogger("CalibrationLogger")
        self.logger.setLevel(logging.INFO)
        formatter = logging.Formatter('%(asctime)s %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
        file_handler = logging.FileHandler(log_file, encoding="utf-8")
        file_handler.setFormatter(formatter)
        # 避免重复添加handler
        if not self.logger.handlers:
            self.logger.addHandler(file_handler)
        self.log_file = log_file

        self.color_map = {
            "left_edge": "r",
            "right_edge": "g",
            "top_edge": "b",
            "bottom_edge": "m"
        }

                # 定义姿态点的颜色和标签
        self.pose_colors = {
            "left_top": "red",      # 左上角点
            "right_top": "blue",    # 右上角点
            "left_bottom": "green", # 左下角点
            "right_bottom": "purple", # 右下角点
            "left_0": "orange",     # 左侧边缘点
            "right_0": "cyan",      # 右侧边缘点
            "top_0": "magenta",     # 上侧边缘点
            "bottom_0": "yellow"    # 下侧边缘点
        }
        
        self.pose_labels = {
            "left_top": "左上角",
            "right_top": "右上角",
            "left_bottom": "左下角",
            "right_bottom": "右下角",
            "left_0": "左侧边缘",
            "right_0": "右侧边缘",
            "top_0": "上侧边缘",
            "bottom_0": "下侧边缘"
        }

    def visualize_pose(self, pose_dict, origin_img):
        """
        可视化姿态点在原始图像上的位置
        
        参数:
            pose_dict: 姿态点字典，格式为 {'point_name': (x, y)}
            origin_img: 原始图像 (BGR格式)
        """
        # 将BGR图像转换为RGB
        img_rgb = cv2.cvtColor(origin_img, cv2.COLOR_BGR2RGB)
        
        # 创建图形
        fig, ax = plt.subplots(figsize=(12, 8))
        
        # 显示原始图像
        ax.imshow(img_rgb)
        ax.set_title("姿态点检测结果")
        
        # 绘制所有姿态点
        for point_name, (x, y) in pose_dict.items():
            color = self.pose_colors.get(point_name, "white")
            label = self.pose_labels.get(point_name, point_name)
            
            # 绘制点
            ax.scatter(x, y, s=100, c=color, edgecolors='black', alpha=0.8)
            
            # 添加标签
            ax.text(x + 10, y + 10, f"{label}", 
                    color=color, fontsize=12, weight='bold',
                    bbox=dict(facecolor='white', alpha=0.7, edgecolor='none'))
        
        # 添加图例
        handles = []
        labels = []
        for point_name in pose_dict.keys():
            if point_name in self.pose_colors:
                color = self.pose_colors[point_name]
                label = self.pose_labels.get(point_name, point_name)
                handles.append(plt.Line2D([0], [0], marker='o', color='w', 
                                          markerfacecolor=color, markersize=10))
                labels.append(label)
        
        ax.legend(handles, labels, loc='upper right', fontsize=10)
        
        plt.tight_layout()
        plt.show()
        
        # 记录日志
        pose_info = "\n".join([f"{self.pose_labels.get(name, name)}: ({x:.1f}, {y:.1f})" 
                              for name, (x, y) in pose_dict.items()])
        self.logger.info(f"姿态点检测结果:\n{pose_info}")
    def visualize_edges(self, edges, edges_depth, origin_img, depth_img):
        msg = (
            f"edges:{edges}\n"
            f"edges_depth:{edges_depth}\n"
        )
        self.logger.info(msg=msg)
        
        # 转为RGB以便matplotlib显示
        origin_img_rgb = cv2.cvtColor(origin_img, cv2.COLOR_BGR2RGB)
        depth_img_vis = cv2.normalize(depth_img, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
        depth_img_vis = cv2.cvtColor(depth_img_vis, cv2.COLOR_GRAY2RGB)

        fig, axs = plt.subplots(2, 2, figsize=(16, 10), gridspec_kw={'height_ratios': [1, 2]})
        
        # 上方：深度分布
        ax_top = axs[0, 0]
        for edge_name, points in edges.items():
            if edge_name not in edges_depth or edges_depth[edge_name] is None:
                continue
            depths = edges_depth[edge_name]
            if not isinstance(points, (list, tuple)) or not isinstance(depths, (list, tuple)):
                continue
            if len(points) != len(depths):
                continue
            ds = [d for d in depths if d is not None]
            ax_top.plot(range(len(ds)), ds, marker='o', 
                    color=self.color_map.get(edge_name, "k"), 
                    label=f"{edge_name} 深度")
        
        ax_top.set_xlabel("采样点序号")
        ax_top.set_ylabel("深度值")
        ax_top.set_title("各边采样点深度分布")
        ax_top.legend()
        ax_top.grid(True)
        
        # 隐藏右上角子图
        axs[0, 1].axis('off')

        # 下方左：origin_img
        ax_left = axs[1, 0]
        ax_left.imshow(origin_img_rgb)
        ax_left.set_title("原图像及采样点深度")
        for edge_name, points in edges.items():
            if edge_name not in edges_depth or edges_depth[edge_name] is None:
                continue
            depths = edges_depth[edge_name]
            for pt, d in zip(points, depths):
                if pt is not None and d is not None and len(pt) == 2:
                    u, v = pt
                    color = self.color_map.get(edge_name, "k")
                    ax_left.scatter(u, v, c=color, s=30)
                    ax_left.text(u, v, str(int(d)), color=color, fontsize=8)
        ax_left.axis('off')

        # 下方右：depth_img
        ax_right = axs[1, 1]
        ax_right.imshow(depth_img_vis)
        ax_right.set_title("深度图及采样点深度")
        for edge_name, points in edges.items():
            if edge_name not in edges_depth or edges_depth[edge_name] is None:
                continue
            depths = edges_depth[edge_name]
            for pt, d in zip(points, depths):
                if pt is not None and d is not None and len(pt) == 2:
                    u, v = pt
                    color = self.color_map.get(edge_name, "k")
                    ax_right.scatter(u, v, c=color, s=30)
                    ax_right.text(u, v, str(int(d)), color=color, fontsize=8)
        ax_right.axis('off')

        plt.tight_layout()
        plt.show()


    def visualize_edges_3d(self,edges_3d):
        """
        可视化3D边缘点。每条边用不同颜色和marker，添加图例和坐标轴标签。
        输入：edges_3d: dict, 形如{"left_edge": [...], ...}
        """
        marker_map = {
            "left_edge": "o",
            "right_edge": "^",
            "top_edge": "s",
            "bottom_edge": "x"
        }
        fig = plt.figure(figsize=(10, 8))
        ax = fig.add_subplot(111, projection='3d')
        for edge_name, points in edges_3d.items():
            if not isinstance(points, (list, tuple)) or len(points) == 0:
                continue
            pts = [pt for pt in points if pt is not None and hasattr(pt, '__len__') and len(pt) == 3]
            if len(pts) == 0:
                continue
            pts = np.array(pts)
            color = self.color_map.get(edge_name, 'k')
            marker = marker_map.get(edge_name, '.')
            ax.scatter(pts[:,0], pts[:,1], pts[:,2], c=color, marker=marker, label=edge_name)
            # 连线
            ax.plot(pts[:,0], pts[:,1], pts[:,2], color=color, alpha=0.5)
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')
        ax.set_title('3D边缘点可视化')
        ax.legend()
        plt.tight_layout()
        plt.show()

    def visualize_fitted_edges(self, edges_3d, fitted_edges):
        """
        可视化原始3D点和拟合后的3D点，以及center点。
        原始点用o，拟合点用三角形，center用*，不同边不同颜色。
        每条边的第一个点和最后一个点用直线连接，并显示长度。
        """
        marker_raw = 'o'
        marker_fit = '^'
        fig = plt.figure(figsize=(10, 8))
        ax = fig.add_subplot(111, projection='3d')
        # 原始点
        for edge_name, points in edges_3d.items():
            pts = [pt for pt in points if pt is not None and hasattr(pt, '__len__') and len(pt) == 3]
            if len(pts) == 0:
                continue
            pts = np.array(pts)
            color = self.color_map.get(edge_name, 'k')
            ax.scatter(pts[:,0], pts[:,1], pts[:,2], c=color, marker=marker_raw, label=f'{edge_name} 原始')
        # 拟合点
        for edge_name, points in fitted_edges.items():
            if edge_name == 'center':
                continue
            pts = [pt for pt in points if pt is not None and hasattr(pt, '__len__') and len(pt) == 3]
            if len(pts) == 0:
                continue
            pts = np.array(pts)
            color = self.color_map.get(edge_name, 'k')
            ax.scatter(pts[:,0], pts[:,1], pts[:,2], c=color, marker=marker_fit, label=f'{edge_name} 拟合')
            # 连接第一个点和最后一个点，并显示长度
            first = pts[0]
            last = pts[-1]
            ax.plot([first[0], last[0]], [first[1], last[1]], [first[2], last[2]], color=color, linestyle='--', linewidth=2)
            # 计算长度
            length = np.linalg.norm(np.array(last) - np.array(first))
            # 在连线中点处显示长度
            mid = (np.array(first) + np.array(last)) / 2
            ax.text(mid[0], mid[1], mid[2], f"{length:.2f}", color=color, fontsize=12, backgroundcolor='w')
        # center点
        center = fitted_edges.get('center', None)
        if center is not None and hasattr(center, '__len__') and len(center) == 3:
            ax.scatter(center[0], center[1], center[2], c='gold', marker='*', s=150, label='center')
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')
        ax.set_title('3D边缘点拟合前后对比')
        ax.legend()
        plt.tight_layout()
        plt.show()

    def visualize_fitted_vector(self, standard_vector, fitted_vector):
        """
        优化可视化方法，解决标签重叠和向量过长问题
        
        参数:
            standard_vector: 标准向量字典
            fitted_vector: 拟合向量字典
        """
        if standard_vector is None:
            standard_vector = fitted_vector
        # 创建3D图形
        fig = plt.figure(figsize=(15, 12))
        ax = fig.add_subplot(111, projection='3d')
        
        # 设置标题
        ax.set_title("Standard vs Fitted Vectors Comparison", fontsize=14)
        
        # 1. 绘制关键点 - 使用不同颜色和标记区分标准与拟合
        self._plot_key_points(ax, standard_vector, 'b', 's', "_std", offset=0.02)  # 蓝色方块表示标准点
        self._plot_key_points(ax, fitted_vector, 'r', 'o', "_fitted", offset=0.02)  # 红色圆圈表示拟合点
        
        # 2. 绘制向量 - 使用更小的长度比例
        self._plot_vectors(ax, standard_vector, color='b', label='Standard', linestyle='-', scale=0.05)
        self._plot_vectors(ax, fitted_vector, color='r', label='Fitted', linestyle='--', scale=0.05)
        
        # 3. 计算并显示距离信息 - 使用不同位置显示避免重叠
        if "center" in fitted_vector and "left_bottom" in fitted_vector:
            center = np.array(fitted_vector["center"])
            left_bottom = np.array(fitted_vector["left_bottom"])
            distance = np.linalg.norm(center - left_bottom)
            
            # 在中心点上方添加距离标注
            ax.text(center[0], center[1], center[2] + 0.05, 
                    f"Fitted Dist: {distance:.4f}", 
                    color='r', fontsize=10, bbox=dict(facecolor='white', alpha=0.7))
        
        if "center" in standard_vector and "left_bottom" in standard_vector:
            center = np.array(standard_vector["center"])
            left_bottom = np.array(standard_vector["left_bottom"])
            distance = np.linalg.norm(center - left_bottom)
            
            # 在中心点下方添加距离标注
            ax.text(center[0], center[1], center[2] - 0.05, 
                    f"Standard Dist: {distance:.4f}", 
                    color='b', fontsize=10, bbox=dict(facecolor='white', alpha=0.7))
        
        # 4. 设置图例 - 只显示主要项目避免混乱
        handles, labels = ax.get_legend_handles_labels()
        unique_labels = []
        unique_handles = []
        for handle, label in zip(handles, labels):
            if label not in unique_labels:
                unique_labels.append(label)
                unique_handles.append(handle)
        ax.legend(unique_handles, unique_labels, loc='best')
        
        # 5. 设置坐标轴标签和范围
        ax.set_xlabel('X', fontsize=12)
        ax.set_ylabel('Y', fontsize=12)
        ax.set_zlabel('Z', fontsize=12)
        
        # 6. 自动调整坐标轴范围
        self._auto_adjust_axis(ax, standard_vector, fitted_vector)
        
        # 7. 设置视角以便更好地观察
        ax.view_init(elev=25, azim=30)
        
        # 8. 添加网格和透明度增强可视化
        ax.grid(True, linestyle='--', alpha=0.7)
        ax.xaxis.pane.fill = False
        ax.yaxis.pane.fill = False
        ax.zaxis.pane.fill = False
        ax.xaxis.pane.set_edgecolor('w')
        ax.yaxis.pane.set_edgecolor('w')
        ax.zaxis.pane.set_edgecolor('w')
        
        # 显示图形
        plt.tight_layout()
        plt.show()

    def _plot_vectors(self, ax, vector_dict, color, label, linestyle, scale=0.1):
        """绘制向量 - 添加比例参数控制长度"""
        import numpy as np
        
        # 需要绘制的边
        edges = ["left_edge", "bottom_edge"]
        
        for edge in edges:
            if edge in vector_dict:
                vec_info = vector_dict[edge]
                start = np.array(vec_info["start_point"])
                unit_vec = np.array(vec_info["unit_vector"])
                
                # 计算箭头的终点（使用比例参数）
                end = start + unit_vec * scale
                
                # 绘制向量
                ax.quiver(start[0], start[1], start[2], 
                        unit_vec[0], unit_vec[1], unit_vec[2],
                        length=scale, color=color, linestyle=linestyle,
                        arrow_length_ratio=0.3, label=label if edge == "left_edge" else "",
                        alpha=0.8)
                
                # 添加标签 - 使用更小的字体和位置偏移
                ax.text(start[0], start[1], start[2], 
                        f"{edge[0]}_s", color=color, fontsize=8)
                ax.text(end[0], end[1], end[2], 
                        f"{edge[0]}_e", color=color, fontsize=8)

    def _plot_key_points(self, ax, vector_dict, color, marker, suffix="", offset=0.01):
        """绘制关键点 - 添加标记类型和位置偏移"""
        import numpy as np
        
        # 关键点列表
        key_points = ["left_bottom", "left_top", "center"]
        
        for point_name in key_points:
            if point_name in vector_dict:
                point = np.array(vector_dict[point_name])
                
                # 绘制点
                ax.scatter(point[0], point[1], point[2], 
                        s=80, marker=marker, color=color, 
                        label=f"{point_name}{suffix}", alpha=0.8)
                
                # 添加标签 - 使用偏移避免重叠
                text_pos = point + np.array([offset, offset, offset])
                ax.text(text_pos[0], text_pos[1], text_pos[2], 
                        f"{point_name}{suffix}", color=color, fontsize=9,
                        bbox=dict(facecolor='white', alpha=0.5))

    def _auto_adjust_axis(self, ax, standard_vector, fitted_vector):
        """自动调整坐标轴范围以获得最佳视图"""
        import numpy as np
        
        # 收集所有点坐标
        all_points = []
        
        # 添加关键点
        for vec_dict in [standard_vector, fitted_vector]:
            for point_name in ["left_bottom", "left_top", "center"]:
                if point_name in vec_dict:
                    all_points.append(np.array(vec_dict[point_name]))
        
        # 添加向量起点终点
        for vec_dict in [standard_vector, fitted_vector]:
            for edge in ["left_edge", "bottom_edge"]:
                if edge in vec_dict:
                    start = np.array(vec_dict[edge]["start_point"])
                    unit_vec = np.array(vec_dict[edge]["unit_vector"])
                    end = start + unit_vec * 0.05  # 使用与可视化相同的比例
                    all_points.append(start)
                    all_points.append(end)
        
        if not all_points:
            return
        
        # 转换为数组并计算范围
        all_points = np.array(all_points)
        min_vals = np.min(all_points, axis=0)
        max_vals = np.max(all_points, axis=0)
        
        # 扩展10%的范围
        range_vals = max_vals - min_vals
        padding = range_vals * 0.1
        min_vals -= padding
        max_vals += padding
        
        # 设置坐标轴范围
        ax.set_xlim(min_vals[0], max_vals[0])
        ax.set_ylim(min_vals[1], max_vals[1])
        ax.set_zlim(min_vals[2], max_vals[2])

    def visualize_vector(self, vector):
        """
        可视化3D向量和中心点
        
        参数:
            vector (dict): 包含以下键:
                - center: 中心点坐标 [x, y, z]
                - bottom_vec: 底部向量 [x, y, z]
                - left_vec: 左侧向量 [x, y, z]
                - z_vec: z轴向量 [x, y, z]
                - bottom_angle: 底部角度(角度值)
                - left_angle: 左侧角度(角度值)
        """
        if vector is None:
            return
        # 创建3D图形
        fig = plt.figure(figsize=(10, 8))
        ax = fig.add_subplot(111, projection='3d')
        
        # 提取数据
        center = np.array(vector['center'])
        bottom_vec = np.array(vector['bottom_vec'])
        left_vec = np.array(vector['left_vec'])
        z_vec = np.array(vector['z_vec'])
        
        # 绘制标准坐标轴（虚线）
        axis_length = 0.2  # 坐标轴长度
        ax.quiver(0, 0, 0, axis_length, 0, 0, color='k', linestyle=':', linewidth=1, arrow_length_ratio=0.1, label='X Axis')
        ax.quiver(0, 0, 0, 0, axis_length, 0, color='k', linestyle=':', linewidth=1, arrow_length_ratio=0.1, label='Y Axis')
        ax.quiver(0, 0, 0, 0, 0, axis_length, color='k', linestyle=':', linewidth=1, arrow_length_ratio=0.1, label='Z Axis')
        
        # 绘制中心点
        ax.scatter(*center, color='r', s=100, label='Center')
        ax.text(*center, f"Center: ({center[0]:.3f}, {center[1]:.3f}, {center[2]:.3f})", 
            color='r', fontsize=10)
        
        # 绘制向量 (从中心点出发)
        vectors = [
            (bottom_vec, 'b', 'Bottom Vector'),
            (left_vec, 'g', 'Left Vector'),
            (z_vec, 'm', 'Z Vector')
        ]
        
        for vec, color, label in vectors:
            ax.quiver(*center, *vec, color=color, length=0.1, normalize=True, 
                    arrow_length_ratio=0.1, label=label)
        
        # 设置坐标轴范围
        min_val = min(center.min(), (center + bottom_vec).min(), 
                    (center + left_vec).min(), (center + z_vec).min()) - 0.1
        max_val = max(center.max(), (center + bottom_vec).max(), 
                    (center + left_vec).max(), (center + z_vec).max()) + 0.1
        
        ax.set_xlim([min_val, max_val])
        ax.set_ylim([min_val, max_val])
        ax.set_zlim([min_val, max_val])
        
        # 添加角度信息（直接显示原始角度值）
        angle_text = (f"Bottom Angle: {vector['bottom_angle']:.2f}°\n"
                    f"Left Angle: {vector['left_angle']:.2f}°")
        
        ax.text2D(0.05, 0.95, angle_text, transform=ax.transAxes,
                bbox=dict(facecolor='white', alpha=0.8))
        
        # 设置图形属性
        ax.set_xlabel('X Axis')
        ax.set_ylabel('Y Axis')
        ax.set_zlabel('Z Axis')
        ax.set_title('3D Vector Visualization')
        
        # 调整图例位置避免重叠
        ax.legend(loc='upper right', bbox_to_anchor=(1.1, 1.1))
        
        # 显示图形
        plt.tight_layout()
        plt.show()

        
